﻿/**************************************************************
 *
 * 唯一标识：163372cd-28a9-4221-80c1-8892d4f810c9
 * 命名空间：Sgr.DistributedLock
 * 创建时间：2024/4/5 21:06:58
 * 机器名称：DESKTOP-HJ4OAG9
 * 创建者：CocoYuan
 * 电子邮箱：fengqinhua2016@163.com
 * 描述：
 *
 **************************************************************/

using Microsoft.Extensions.Logging;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace Sgr.DistributedLock
{
    public class LocalDistributedLock : IDistributedLock, IDisposable
    {
        private readonly ILogger _logger;
        private readonly ConcurrentDictionary<string, SemaphoreRef> _semaphores = new();
        private readonly int _maxLocks;
        private readonly TimeSpan _lockTimeout;
        private volatile bool _disposed;

        public LocalDistributedLock(ILogger<LocalDistributedLock> logger)
        {
            _logger = logger;
            _maxLocks = 1024; // 默认最大锁数量
            _lockTimeout = TimeSpan.FromSeconds(30); // 默认锁超时时间
        }

        private class SemaphoreRef : IDisposable
        {
            public SemaphoreSlim Semaphore { get; }
            private int _refCount;
            private readonly object _syncLock = new();
            public bool IsDisposed { get; private set; }

            public SemaphoreRef()
            {
                Semaphore = new SemaphoreSlim(1, 1);
            }

            public int AddRef()
            {
                return Interlocked.Increment(ref _refCount);
            }

            public int Release()
            {
                int newCount = Interlocked.Decrement(ref _refCount);
                if (newCount < 0)
                {
                    Interlocked.Increment(ref _refCount);
                    throw new InvalidOperationException("引用计数不能为负数");
                }
                return newCount;
            }

            public int RefCount => Volatile.Read(ref _refCount);

            public void Dispose()
            {
                lock (_syncLock)
                {
                    if (!IsDisposed)
                    {
                        Semaphore.Dispose();
                        IsDisposed = true;
                    }
                }
            }
        }

        public ILocker AcquireLock(string key, TimeSpan? expiration = null)
        {
            CheckDisposed();
            Check.StringNotNullOrEmpty(key, nameof(key));

            if (expiration != null && expiration < TimeSpan.Zero)
                throw new ArgumentException("Expiration must be greater than or equal to zero.", nameof(expiration));

            var semaphoreRef = GetOrCreateLock(key);//add or update,use + 1

            try
            {
                if (!semaphoreRef.Semaphore.Wait(_lockTimeout))         //wait for the Semaphore
                    throw new TimeoutException($"获取分布式锁超时: {key}");

                return new Locker(this, semaphoreRef, key, expiration); //release the Semaphore , ReleaseLock
            }
            catch
            {
                ReleaseLock(key, semaphoreRef); //use -1 , if refcount = 0,remove the key
                throw;
            }
        }

        public async Task<ILocker> AcquireLockAsync(string key, TimeSpan? expiration = null)
        {
            CheckDisposed();
            Check.StringNotNullOrEmpty(key, nameof(key));

            if (expiration != null && expiration < TimeSpan.Zero)
                throw new ArgumentException("Expiration must be greater than or equal to zero.", nameof(expiration));

            var semaphoreRef = GetOrCreateLock(key);

            try
            {
                if (!await semaphoreRef.Semaphore.WaitAsync(_lockTimeout))
                    throw new TimeoutException($"获取分布式锁超时: {key}");

                return new Locker(this, semaphoreRef, key, expiration);
            }
            catch
            {
                ReleaseLock(key, semaphoreRef);
                throw;
            }
        }

        public ILocker? TryAcquireLock(string key, TimeSpan timeout, TimeSpan? expiration = null)
        {
            CheckDisposed();
            Check.StringNotNullOrEmpty(key, nameof(key));

            if (timeout < TimeSpan.Zero)
                throw new ArgumentException("Timeout must be greater than or equal to zero.", nameof(timeout));

            if (expiration != null && expiration < TimeSpan.Zero)
                throw new ArgumentException("Expiration must be greater than or equal to zero.", nameof(expiration));

            var semaphoreRef = GetOrCreateLock(key);

            try
            {
                if (semaphoreRef.Semaphore.Wait(timeout))
                    return new Locker(this, semaphoreRef, key, expiration);

                ReleaseLock(key, semaphoreRef);
                return null;
            }
            catch
            {
                ReleaseLock(key, semaphoreRef);
                throw;
            }
        }

        public async Task<ILocker?> TryAcquireLockAsync(string key, TimeSpan timeout, TimeSpan? expiration = null)
        {
            CheckDisposed();
            Check.StringNotNullOrEmpty(key, nameof(key));

            if (timeout < TimeSpan.Zero)
                throw new ArgumentException("Timeout must be greater than or equal to zero.", nameof(timeout));

            if (expiration != null && expiration < TimeSpan.Zero)
                throw new ArgumentException("Expiration must be greater than or equal to zero.", nameof(expiration));

            var semaphoreRef = GetOrCreateLock(key);

            try
            {
                if (await semaphoreRef.Semaphore.WaitAsync(timeout))
                    return new Locker(this, semaphoreRef, key, expiration);

                ReleaseLock(key, semaphoreRef);
                return null;
            }
            catch
            {
                ReleaseLock(key, semaphoreRef);
                throw;
            }
        }

        private SemaphoreRef GetOrCreateLock(string key)
        {
            if (_semaphores.Count >= _maxLocks)
            {
                throw new InvalidOperationException($"已达到最大锁数量限制({_maxLocks})，请稍后重试");
            }

            return _semaphores.AddOrUpdate(
                key,
                _ =>
                {
                    var newRef = new SemaphoreRef();
                    newRef.AddRef();
                    return newRef;
                },
                (_, existingRef) =>
                {
                    existingRef.AddRef();
                    return existingRef;
                });
        }

        private void ReleaseLock(string key, SemaphoreRef semaphoreRef)
        {
            int newCount = semaphoreRef.Release();
            if (newCount == 0)
            {
                if (_semaphores.TryRemove(key, out var lockRef))
                {
                    try
                    {
                        if (lockRef.RefCount == 0)
                        {
                            lockRef.Dispose();
                        }
                        else
                        {
                            _semaphores.TryAdd(key, lockRef);
                        }
                    }
                    catch (Exception ex)
                    {
                        _logger.LogError(ex, "清理键 {Key} 的锁时发生错误", key);
                        if (!lockRef.IsDisposed)
                        {
                            _semaphores.TryAdd(key, lockRef);
                        }
                    }
                }
            }
        }

        private class Locker : ILocker
        {
            private readonly LocalDistributedLock _localLock;
            private readonly SemaphoreRef _semaphoreRef;
            private readonly string _key;
            private readonly CancellationTokenSource? _cts;
            private volatile int _released;
            private volatile bool _disposed;

            public Locker(LocalDistributedLock localLock, SemaphoreRef semaphoreRef, string key, TimeSpan? expiration)
            {
                _localLock = localLock;
                _semaphoreRef = semaphoreRef;
                _key = key;

                if (expiration.HasValue && expiration.Value != TimeSpan.MaxValue)
                {
                    _cts = new CancellationTokenSource(expiration.Value);
                    _cts.Token.Register(Release);
                }
            }

            private void Release()
            {
                if (Interlocked.Exchange(ref _released, 1) == 0)
                {
                    _semaphoreRef.Semaphore.Release();
                    _localLock.ReleaseLock(_key, _semaphoreRef);
                }
            }

            public ValueTask DisposeAsync()
            {
                Dispose();
                return default;
            }

            public void Dispose()
            {
                if (_disposed)
                    return;

                _disposed = true;
                _cts?.Dispose();
                Release();
            }
        }

        private void CheckDisposed()
        {
            ObjectDisposedException.ThrowIf(_disposed, this);
        }

        public void Dispose()
        {
            if (_disposed)
                return;

            _disposed = true;

            foreach (var pair in _semaphores)
            {
                pair.Value.Dispose();
            }
            _semaphores.Clear();

            GC.SuppressFinalize(this);
        }

        public bool IsLockAcquired(string key)
        {
            CheckDisposed();
            Check.NotNull(key, nameof(key));

            if (!_semaphores.TryGetValue(key, out var semaphoreRef))
                return false;

            // 检查信号量是否已释放
            if (semaphoreRef.IsDisposed)
                return false;

            try
            {
                // 检查当前信号量计数
                return semaphoreRef.Semaphore.CurrentCount == 0;
            }
            catch (ObjectDisposedException)
            {
                // 如果在检查过程中信号量被释放，返回false
                return false;
            }
        }

        public Task<bool> IsLockAcquiredAsync(string key)
        {
            return Task.FromResult(IsLockAcquired(key));
        }
    }
}